#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : microbiome_stacked_bar.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/10/21 14:35
@Description :
分组绘制微生态物种百分比堆叠图
使用人工精选的高对比离散调色板（按平均丰度排序）
堆叠顺序与图例一致
"""
from argparse import ArgumentParser, Namespace
from pathlib import Path

import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

# ✅ 人工精选高对比离散调色板（40 种）
HIGH_CONTRAST_COLORS = [
    "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd",
    "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf",
    "#393b79", "#637939", "#8c6d31", "#843c39", "#7b4173",
    "#3182bd", "#e6550d", "#31a354", "#756bb1", "#636363",
    "#9ecae1", "#fdae6b", "#a1d99b", "#bcbddc", "#969696",
    "#6baed6", "#fd8d3c", "#74c476", "#9e9ac8", "#bdbdbd",
    "#ff0000", "#00ccff", "#ff66cc", "#33cc33", "#cc9933",
    "#3399ff", "#cc3366", "#9999ff", "#ffcc00", "#66cccc"
]

def _read_infiles(samples, infiles, outfile) -> pd.DataFrame:
    merged_df = pd.DataFrame()
    for sample, infile in zip(samples, infiles):
        df = pd.read_csv(infile, sep="\t", usecols=["name", "new_est_reads"])
        df.rename(columns={"new_est_reads": sample}, inplace=True)
        df[sample] = df[sample] / df[sample].sum()
        merged_df = df if merged_df.empty else merged_df.merge(df, on="name", how="outer")

    merged_df.fillna(0, inplace=True)
    Path(outfile).parent.mkdir(parents=True, exist_ok=True)
    merged_df.to_csv(outfile, sep="\t", index=False)
    return merged_df.melt(id_vars="name", var_name="Sample", value_name="abundance")

def _get_discrete_colors(num_species: int):
    if num_species <= len(HIGH_CONTRAST_COLORS):
        return HIGH_CONTRAST_COLORS[:num_species]
    else:
        repeats = (num_species // len(HIGH_CONTRAST_COLORS)) + 1
        return (HIGH_CONTRAST_COLORS * repeats)[:num_species]

def _plot_stacked_bar(in_df, out_pdf, out_png, top=15, width=32, height=12, dpi=360):
    groups = in_df["Group"].unique()
    fig, axes = plt.subplots(2, 2, figsize=(width, height), sharey=True)
    axes = axes.flatten()

    for ax, group in zip(axes, groups):
        group_data = in_df[in_df["Group"] == group]

        pivot_df = group_data.pivot_table(
            index="Sample", columns="name", values="abundance", fill_value=0
        )

        # 按平均丰度排序
        mean_abundance = pivot_df.mean(axis=0).sort_values(ascending=False)
        top_species = mean_abundance.head(top).index.tolist()
        other_species = [sp for sp in pivot_df.columns if sp not in top_species]

        plot_df = pivot_df.copy()
        if other_species:
            plot_df["Other"] = plot_df[other_species].sum(axis=1)
            plot_df = plot_df.drop(columns=other_species)

        # 高丰度 → 低丰度 + Other
        species_order = sorted([sp for sp in plot_df.columns if sp != "Other"],
                               key=lambda sp: mean_abundance.get(sp, 0), reverse=True)
        if "Other" in plot_df.columns:
            species_order.append("Other")
        plot_df = plot_df[species_order]

        colors = _get_discrete_colors(len(species_order))

        # 计算累计高度
        cumulative = np.zeros(len(plot_df))

        # 先计算每个条形的顶部位置
        totals = plot_df.sum(axis=1).values  # 每个样本总高度

        # 从上到下绘制
        for i, sp in enumerate(plot_df.columns):
            # 当前柱子的 bottom = 总高度 - 当前值 - 已绘制高度
            bottom = totals - cumulative - plot_df[sp].values
            ax.bar(plot_df.index, plot_df[sp], bottom=bottom, color=colors[i], width=0.8)
            ax.set_xticks(np.arange(len(plot_df.index)))
            ax.set_xticklabels(plot_df.index, rotation=90, ha='center')
            cumulative += plot_df[sp].values

        # ✅ 图例从上到下和柱子视觉顺序一致
        ax.legend(plot_df.columns, fontsize=9, frameon=False,
                  prop={"style":"italic"}, bbox_to_anchor=(1.02,1), loc="upper left")

        ax.set_title(group, fontsize=14)
        ax.set_ylabel("Relative abundance", fontsize=12)
        ax.set_xlabel("")
        ax.set_xticklabels(plot_df.index, rotation=90, ha='center')

    for i in range(len(groups), 4):
        fig.delaxes(axes[i])

    plt.tight_layout(rect=(0.05, 0.05, 0.95, 0.95))
    Path(out_pdf).parent.mkdir(parents=True, exist_ok=True)
    Path(out_png).parent.mkdir(parents=True, exist_ok=True)
    plt.savefig(out_png, dpi=dpi)
    plt.savefig(out_pdf)
    plt.close()

def main(args: Namespace):
    mpl.use("cairo")
    group_df = pd.read_csv(args.groupfile, sep="\t", encoding="utf-8")
    merged_df = _read_infiles(
        args.samples.strip().split(","),
        args.infiles.strip().split(","),
        args.outfile,
    )
    merged_df = merged_df.merge(group_df, on="Sample", how="left")
    _plot_stacked_bar(
        in_df=merged_df,
        out_pdf=args.outpdf,
        out_png=args.outpng,
        top=args.top,
        width=args.width,
        height=args.height,
        dpi=args.dpi,
    )

if __name__ == "__main__":
    parser = ArgumentParser(description="分组绘制微生态物种相对丰度堆叠柱状图（按平均丰度排序）")
    parser.add_argument("--samples", type=str, required=True, help="样本名称列表（逗号分隔）")
    parser.add_argument("--infiles", type=str, required=True, help="每个样本对应的丰度文件（逗号分隔）")
    parser.add_argument("--groupfile", type=str, required=True, help="样本分组文件")
    parser.add_argument("--outfile", type=str, required=True, help="输出合并表路径")
    parser.add_argument("--outpdf", type=str, required=True, help="输出 PDF 文件路径")
    parser.add_argument("--outpng", type=str, required=True, help="输出 PNG 文件路径")
    parser.add_argument("--top", type=int, default=15, help="每组显示的 Top 物种数量")
    parser.add_argument("--width", type=float, default=32, help="图像宽度")
    parser.add_argument("--height", type=float, default=12, help="图像高度")
    parser.add_argument("--dpi", type=float, default=360, help="输出分辨率")
    main(parser.parse_args())
